import time
import asyncio
from .stream_server import StreamServer
from asyncio import StreamReader, StreamWriter
from .handle import RpcHandle, Conn
from .route import Route
from lib.log import Log
from typing import Dict, List, Optional
from . import ServiceDiscover, DiscoverOption, Service


class RpcServer(StreamServer, Service):
    """
    rpc服务端
    """

    def __init__(self, ip: str, port: int, discover_option: DiscoverOption = None):
        super().__init__(ip, port, name=type(self).__name__)
        self.address: str = f"{ip}:{port}"
        self._name: str = f"worker_{self.address}"
        self.address2Conn: Dict[str, RpcHandle] = {}  # client节点address对应的conn
        self.tag2ConnList: Dict[str, List[RpcHandle]] = {}  # client节点别名对应的conn
        self.route: Route = Route()  # 服务路由
        self.init_route()  # 路由初始化
        self.is_close = False
        self.log = Log(f"worker_{ip}_{port}")

        self.discover_node_name: str = ""
        self.discover_heartbeat_time: int = 0
        self.discover_instance: Optional[ServiceDiscover] = None
        if discover_option:
            self.discover_node_name = discover_option.node_name
            self.discover_heartbeat_time = discover_option.heartbeat_time
            self.discover_instance = discover_option.instance

    @property
    def name(self) -> str:
        return self._name

    def init_route(self) -> None:
        """初始化路由"""
        _key2Func = {
            "register": self.register,
        }
        for key, func in _key2Func.items():
            self.route.add(key, func)

    async def run(self) -> None:
        await self.register_service()
        await super().run()

    async def close(self):
        """服务关闭"""
        self.is_close = True
        await self.unregister_service()
        self.main_server.close()
        await self.main_server.wait_closed()

    async def heartbeat_active(self):
        """心跳"""
        while not self.is_close:
            await asyncio.sleep(self.discover_heartbeat_time - 2)
            # 更新心跳时间
            await self.discover_instance.register_service(self.discover_node_name, self.address, int(time.time()))

    async def register_service(self):
        """服务注册"""
        if not self.discover_instance:
            return

        # 注册
        await self.discover_instance.register_service(self.discover_node_name, self.address, int(time.time()))
        # 通知新的service到来
        await self.discover_instance.notice(self.discover_node_name, f"address={self.address}&act=add")
        # 心跳保活
        asyncio.create_task(self.heartbeat_active())

    async def unregister_service(self):
        """服务注销"""
        if self.discover_instance:
            # 销毁
            await self.discover_instance.remove_service(self.discover_node_name, self.address)
            # 通知service销毁
            await self.discover_instance.notice(self.discover_node_name, f"address={self.address}&act=del")

    async def handle(self, reader: StreamReader, writer: StreamWriter) -> None:
        """消息处理"""
        rpc_handle = RpcHandle(self, Conn(reader, writer, name=self.name), self.route, close_callback=self.client_close)
        rpc_handle.set_log(self.log)
        await rpc_handle.run()

    async def client_close(self, conn: RpcHandle) -> None:
        """客户端连接断开"""
        del_address = None
        for address, _conn in self.address2Conn.items():
            if _conn == conn:
                del_address = address
                break
        if del_address:
            del self.address2Conn[del_address]

        for _, connList in self.tag2ConnList.items():
            if conn in connList:
                connList.remove(conn)

    async def register(self, conn: RpcHandle, address: str, tag: str = None) -> str:
        """client注册"""
        if address in self.address2Conn:
            self.log.warning(f'【{self.address}】 already register! from->[{address}]')

        self.address2Conn[address] = conn

        # 别名
        if tag:
            self.tag2ConnList.setdefault(tag, []).append(conn)

        return f"register {self.address} success!"


if __name__ == '__main__':
    server = RpcServer("10.0.0.153", 9001)

    try:
        asyncio.run(server.run())
    except KeyboardInterrupt:
        pass
